package site.codeyin.javautils.utils.filedownLoad;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import site.codeyin.javautils.constant.BizConstant;

import javax.annotation.Resource;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;

@Slf4j
@Service
public class MultiThreadedDownloader {

    @Resource(name = "downloadExecutorService")
    private ThreadPoolExecutor downloadExecutorService;

    private static final int FILE_PART_SIZE = 15; // 文件分片的数量
    private static int flag = 0; // 下载成功的标志

    public synchronized Boolean downloadFile(String serverPath, String localFilePath) {

        long startTime = System.currentTimeMillis();

        try {
            URL url = new URL(serverPath);
            HttpURLConnection connection = (HttpURLConnection) url.openConnection();
            connection.setRequestMethod("HEAD");
            int fileSize = connection.getContentLength();
            connection.disconnect();

            log.info("File size: {} bytes", fileSize);

            if (fileSize <= 0) {
                log.error("Invalid file size");
                return false;
            }

            File file = new File(localFilePath).getParentFile();
            if (!file.exists()) {
                boolean mkdir = file.mkdirs();
                log.info("Create dir:{}", mkdir);
            }

            RandomAccessFile raf = new RandomAccessFile(localFilePath, "rwd");
            raf.setLength(fileSize);
            raf.close();

            // 计算每个线程需要下载的文件大小
            long partSize = (long) Math.ceil((double) fileSize / FILE_PART_SIZE);

            CountDownLatch latch = new CountDownLatch(FILE_PART_SIZE);

            for (int i = 0; i < FILE_PART_SIZE; i++) {
                long start = i * partSize;
                long end = (i == FILE_PART_SIZE - 1) ? fileSize - 1 : start + partSize - 1;
                downloadExecutorService.submit(new DownloadTask(String.valueOf(i), serverPath, localFilePath, start, end, latch));
            }

            if (flag == 1) {
                log.error("Download failed");
                return false;
            }

            latch.await();
            // 记录下载文件花费的时间
            log.info("Download completed in {} ms  rate:{}MB/s", System.currentTimeMillis() - startTime, DataRateCalculator.calculateDataRate(fileSize, System.currentTimeMillis() - startTime));
            return true;
        } catch (IOException | InterruptedException e) {
            log.error("error during file download: {}", e.getMessage());
            return false;
        } finally {
            flag = 0;
        }
    }

    static class DownloadTask implements Runnable {
        private final String serverPath;
        private final String localPath;
        private final String threadId;
        private final long start;
        private final long end;
        private final CountDownLatch latch;

        public DownloadTask(String threadId, String serverPath, String localPath, long start, long end, CountDownLatch latch) {
            this.threadId = threadId;
            this.serverPath = serverPath;
            this.localPath = localPath;
            this.start = start;
            this.end = end;
            this.latch = latch;
        }

        @Override
        public void run() {
            log.info("Thread: {} started downloading part from {} to {}", threadId, start, end);

            RandomAccessFile raf = null;
            HttpURLConnection connection = null;
            int bufferSize = BizConstant.BUFFER_SIZE;
            try {
                long startTime = System.currentTimeMillis();
                raf = new RandomAccessFile(localPath, "rwd");
                connection = (HttpURLConnection) new URL(serverPath).openConnection();

                connection.setRequestProperty("Connection", "keep-alive");
                connection.setRequestProperty("Range", "bytes=" + start + "-" + end);
                connection.setConnectTimeout(5000);
                connection.setRequestMethod("GET");


                BufferedInputStream bis = new BufferedInputStream(connection.getInputStream());

                ByteArrayOutputStream baos = new ByteArrayOutputStream();

                byte[] buffer = new byte[bufferSize];
                int bytesRead;
                long totalBytesRead = 0;

                while ((bytesRead = bis.read(buffer, 0, (int) Math.min(bufferSize, (end - start - totalBytesRead + 1)))) != -1) {
                    baos.write(buffer, 0, bytesRead);
                    totalBytesRead += bytesRead;

                    if (totalBytesRead >= (end - start + 1)) {
                        break;
                    }
                }

                byte[] partData = baos.toByteArray();

                raf.seek(start);
                raf.write(partData);
                log.info("Thread: {} downloaded {} bytes finished for part from {} to {}, rate:{}MB/s", threadId, partData.length, start, end, DataRateCalculator.calculateDataRate(partData.length, System.currentTimeMillis() - startTime));
            } catch (IOException e) {
                flag = 1;
                log.error("Thread:{}error during file download: {}", threadId, e.getMessage());
            } finally {
                try {
                    assert raf != null;
                    assert connection != null;
                    raf.close();
                    connection.disconnect();
                } catch (IOException e) {
                    log.error("Thread:{}error during file download: {}", threadId, e.getMessage());
                }
                latch.countDown();
            }
        }
    }

    public static void main(String[] args) {
        String serverPath = "https://boot-video.xuexi.cn/video/1004/p/fc717e16c2ea5deb638713c3f88a136d-7721b113f8a14dd38b6aeafb14d9399a-3.mp4";


        MultiThreadedDownloader downloader = new MultiThreadedDownloader();
        downloader.downloadFile(serverPath, BizConstant.LOCAL_FILE_PATH + new File(serverPath).getName());
    }
}